import torch
import numpy as np
from torch.utils.data import DataLoader
import json
import matplotlib.pyplot as plt
import torch.serialization

from configs.prototype_config import PrototypeConfig
from utils.data_utils import load_dataset
from models.prototype import PrototypeEnhancedCNNTransformer
from utils.metrics import evaluate_model, plot_confusion_matrix


def convert_to_serializable(obj):
    """将NumPy数据类型转换为JSON可序列化的Python原生类型"""
    if isinstance(obj, np.integer):
        return int(obj)
    elif isinstance(obj, np.floating):
        return float(obj)
    elif isinstance(obj, np.ndarray):
        return obj.tolist()
    elif isinstance(obj, dict):
        return {k: convert_to_serializable(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [convert_to_serializable(v) for v in obj]
    else:
        return obj


def load_model_weights(model, checkpoint_path):
    """安全加载模型权重"""
    try:
        # 首先尝试使用weights_only=True加载
        checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=True)
        model.load_state_dict(checkpoint['model_state_dict'])
        return model
    except:
        # 如果失败，尝试使用weights_only=False加载
        print("Warning: Using weights_only=False to load model. Only do this if you trust the source.")
        checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
        model.load_state_dict(checkpoint['model_state_dict'])
        return model


def main():
    # 初始化配置
    config = PrototypeConfig()

    # 加载数据
    _, test_dataset = load_dataset(config)
    test_loader = DataLoader(
        test_dataset,
        batch_size=config.batch_size,
        pin_memory=True if config.device == "cuda" else False
    )

    # 加载模型
    model = PrototypeEnhancedCNNTransformer(config).to(config.device)

    # 修改后的模型加载方式
    try:
        # 尝试使用安全方式加载
        checkpoint = torch.load("results/best_model.pth", map_location=config.device, weights_only=True)
    except:
        # 如果失败，使用非安全方式加载（仅当你信任模型来源时）
        print("Warning: Falling back to weights_only=False loading. Only do this if you trust the model source.")
        checkpoint = torch.load("results/best_model.pth", map_location=config.device, weights_only=False)

    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()

    # 类别名称
    class_names = [
        'lying', 'sitting', 'run', 'walk', 'pickup',
        'wave', 'jump', 'squat', 'sitdown', 'standup', 'fall'
    ]

    # 评估模型
    results = evaluate_model(model, test_loader, config.device, class_names)

    # 计算fall类别的准确率
    fall_mask = np.array(results['true_labels']) == 10
    fall_acc = np.mean(np.array(results['predictions'])[fall_mask] ==
                       np.array(results['true_labels'])[fall_mask]) if any(fall_mask) else 0.0

    # 打印结果
    print(f"\n测试准确率: {results['accuracy']:.4f}")
    print(f"Fall类别准确率: {fall_acc:.4f}")
    print("\n分类报告:")
    print(results['classification_report_str'])

    # 准备要保存的数据
    results_to_save = {
        'accuracy': convert_to_serializable(results['accuracy']),
        'fall_accuracy': convert_to_serializable(fall_acc),
        'confusion_matrix': convert_to_serializable(results['confusion_matrix']),
        'classification_report': convert_to_serializable(results['classification_report_dict']),
        'true_labels': convert_to_serializable(results['true_labels']),
        'predictions': convert_to_serializable(results['predictions'])
    }

    # 保存结果
    with open('results/evaluation_results.json', 'w') as f:
        json.dump(results_to_save, f, indent=2)

    # 绘制混淆矩阵
    plot_confusion_matrix(results['confusion_matrix'], class_names)

    print("\n评估完成，结果已保存到results目录")


if __name__ == "__main__":
    main()